import os
import numpy as np
import SimpleITK as sitk


def delete_unused_nii_files(nii_folder_path):
    """删除指定路径下的，与病情诊断无关的 nii 文件

    :param nii_folder_path: 待处理的nii的存储路径
    """
    for nii_name in os.listdir(nii_folder_path):
        if ('RESP' in nii_name) or ('ep2d' in nii_name) or ('in-opp' in nii_name) \
                or ('cor' in nii_name) or ('test' in nii_name) or ('fs' not in nii_name):
            if ('T2WI' in nii_name) or ('Pre' in nii_name) or ('Arterial' in nii_name) \
                    or ('Venous' in nii_name) or ('Delay' in nii_name):
                break  # 不删除已经改过名字的nii（防止误删）
            delete_file_path = os.path.join(nii_folder_path, nii_name)
            os.remove(delete_file_path)
            print("Delete: ", delete_file_path)


def rename_nii_files_by_phases(nii_folder_path):
    """根据序列时期重命名nii文件

    :param nii_folder_path:待处理的nii的存储路径
    :return:
    """
    nii_file_list = os.listdir(nii_folder_path)
    if len(nii_file_list) != 5:
        print("请根据相关文档，手动删除多余的nii文件，保留有效文件\n"
              "目标路径:", nii_folder_path)
        return
    if nii_file_list[0].find('_') == -1:
        print("该路径下文件不需要重命名\n目标路径:", nii_folder_path)
        return
    patient_info = nii_folder_path.split('\\')[-1]
    patient_name = patient_info.split('_')[0]  # 根据敏感度需要隐藏掉患者的姓名
    study_id = patient_info.split('_')[1]
    print(nii_file_list)

    nii_number_list = []
    for nii_name in nii_file_list:
        nii_number_list.append(int(nii_name.split('-')[0]))
    print(nii_number_list)

    phase = ('T2WI', 'Pre', 'Arterial', 'Venous', 'Delay')
    import sys
    max = sys.maxsize
    for i in range(5):
        min_index = np.argmin(nii_number_list)
        min_index_filename = nii_file_list[min_index]
        old_file_path = os.path.join(nii_folder_path, min_index_filename)
        new_file_path = os.path.join(nii_folder_path, f"{patient_name}-{study_id}-{i}-{phase[i]}.nii")
        os.rename(old_file_path, new_file_path)
        print(f"Renamed {old_file_path} to {new_file_path}")
        nii_number_list[min_index] = max


def nii_resize(image_path, save_folder="", standard_size=0):
    """
    将256*208或320*260的nii图像填充边界(补0),使其大小变为320*320
    图像要求：长宽尺寸均要小于等于320，且为偶数

    Args:
        image_path (string): 待修改大小的nii文件存放路径
        save_folder (string): 修改大小后的nii文件存放路径
        standard_size (int): 修改后标准图像的边长
    Returns:
        图像尺寸符合要求且文件修改成功则返回若干nii文件
    """
    if save_folder != "" and not os.path.exists(save_folder):
        os.mkdir(save_folder)  # 如果保存文件夹不存在，则创建

    for file in os.listdir(image_path):  # 遍历图像列表
        if file.endswith(".nii"):
            print('-' * 50, "\n正在处理: ", file, sep='')
            global img_itk
            try:
                img_itk = sitk.ReadImage(os.path.join(image_path, file))  # 读取医学图像
            except Exception as e:
                print('ReadImage Error:', e)
                continue

            original_size = img_itk.GetSize()  # 获取原始图像尺寸和间距
            image_array = sitk.GetArrayFromImage(img_itk)  # 获取图像数据

            need_standard_size = standard_size if standard_size != 0 else original_size[0]  # 获取设定的标准值
            pad_height = (need_standard_size - original_size[1]) // 2  # 计算填充量
            pad_weight = (need_standard_size - original_size[0]) // 2
            print("before_size: ", original_size, standard_size, pad_height, pad_weight)
            global padded_image_array
            try:
                padded_image_array = np.pad(image_array, (
                    (0, 0), (pad_height, pad_height), (pad_weight, pad_weight)
                ), mode='constant', constant_values=0)  # 使用np.pad对图像进行上下填充
            except Exception as e:
                print('padded Error:', e)
                continue

            resized_image = sitk.GetImageFromArray(padded_image_array)  # 创建新的 SimpleITK 图像

            resized_image.SetOrigin(img_itk.GetOrigin())  # 将元数据设置为与原始图像相同
            resized_image.SetSpacing(img_itk.GetSpacing())
            resized_image.SetDirection(img_itk.GetDirection())

            file_name = os.path.basename(file)
            save_path = os.path.join(save_folder, file_name)  # 生成新的文件名并保存图像
            print("before_size: ", original_size, "\nafter_size : ", resized_image.GetSize(), sep='')
            print("file_name:", save_path)

            sitk.WriteImage(resized_image, save_path)  # 保存调整大小后的 NIfTI 图像
